# -*- coding: UTF-8 -*-

from llm.llm import llm

# from tools.invoke_api import *
# from trace_agent.trace_agent import *
# from utils.utils import *


from langchain.prompts import (
    SystemMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)


summary = """
"header": {
    "title": "Ubuntu资产遭遇高危反弹Shell攻击",
    "desc": "taomingwei-node3-ubuntu自2022年4月3日至2025年4月8日共遭受13次网络攻击，包括检测到Linux可疑命令和高风险的perl反弹shell事件。建议立即排查并采取措施防止进一步恶意活动。"
}
"""


def main():
    sys_tmpl = """
    请将我输入的内容转换成结构化的 HTML + CSS + JavaScript 代码，要求：

    设计风格：现代、美观、专业，使用卡片式布局、平滑动画和精心设计的排版。

    样式要求：
    不要对body、div、p等内置标签初始化，系统已经初始化了
    使用 Flexbox/Grid 优化布局

    代码结构：严格按以下格式返回：
    ```html
    <style>
    /* 你的 CSS 代码 */
    </style>

    <body>
    <!-- 你的 HTML 代码 -->
    </body>

    <script>
    /* 你的 JavaScript 代码 */
    </script>
    ```
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    # print("chat_prompt", chat_prompt.content)

    chain = chat_prompt | llm

    res = chain.invoke({"summary": summary})

    print("res", res)

    with open("test.html", "w") as f:
        f.write(res.text())

    # 保存结果

    print(res.text())


main()
